
#include "token_shift_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BUFFER_NUM = 1;
const uint32_t BLOCK_SIZE = 32;
const int32_t DATATYPE2 = 4;
const int32_t BYTESIZE = 2;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    TokenShiftTilingData tiling;
    const gert::StorageShape* x_shape = context->GetInputShape(0);

    uint32_t B_length = x_shape->GetStorageShape().GetDim(0);
    uint32_t T_length = x_shape->GetStorageShape().GetDim(1);
    uint32_t C_length = x_shape->GetStorageShape().GetDim(2);

    uint32_t sizeofdatatype;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    // std::cout << "SocVersion:" << SocVersion << std::endl;
    // printf("socVersion:%s\n", socVersion);
    // if(socVersion == platform_ascendc::SocVersion::ASCEND310P) {
    //     std::cout << "ASCEND310P:" << std::endl;
    // }
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNum();
    auto dt = context->GetInputTensor(0)->GetDataType();
    if (dt == ge::DT_FLOAT16)
    {
        sizeofdatatype = DATATYPE2;
    }
    uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
    int32_t NUM = 20;
    if(B_length > 1) {
        context->SetTilingKey(100);
        NUM = 20;
    } else {
        NUM = 12;
        context->SetTilingKey(200);
    }
    uint32_t block_num = ((ub_size) / BLOCK_SIZE / 2) / NUM;
    block_num = block_num <= BYTESIZE ? block_num : block_num / BYTESIZE * BYTESIZE;
    uint32_t tiling_size = block_num * ALIGN_NUM;

    aivNum = aivNum >= B_length ? B_length : aivNum;
    uint32_t core_rows = B_length / aivNum;
    uint32_t core_rows_remain = B_length % aivNum;
    uint32_t X_length_align = (B_length * T_length * C_length + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

    tiling.set_B_length(B_length);
    tiling.set_T_length(T_length);
    tiling.set_C_length(C_length);
    tiling.set_tiling_size(tiling_size);
    tiling.set_core_rows(core_rows);
    tiling.set_core_rows_remain(core_rows_remain);
    tiling.set_X_length_align(X_length_align);

    if(B_length == 1) {
        aivNum = ascendcPlatform.GetCoreNum();
        aivNum = aivNum >= T_length ? T_length : aivNum;
        uint32_t core_rows = T_length / aivNum;
        uint32_t core_rows_remain = T_length % aivNum;
        tiling.set_core_rows(core_rows);
        tiling.set_core_rows_remain(core_rows_remain);

        if(C_length <= tiling_size && (C_length % ALIGN_NUM == 0)) {
            context->SetTilingKey(201);
            tiling_size = C_length;
            tiling.set_tiling_size(tiling_size);
        }
        if(T_length == 1) {
            if(C_length % ALIGN_NUM == 0){
                context->SetTilingKey(202);
                tiling_size = C_length;
                tiling.set_tiling_size(tiling_size);
                aivNum = 1;
                if(socVersion == platform_ascendc::SocVersion::ASCEND310P) {
                    std::cout << "ASCEND310P:" << std::endl;
                    context->SetTilingKey(204);
                    aivNum = 6;
                }
            }
        }
    } else {
        // aivNum = ascendcPlatform.GetCoreNum();
        if(B_length <= aivNum) {
            if(C_length < tiling_size) {
                aivNum = 1;
                uint32_t core_rows = B_length;
                uint32_t core_rows_remain = 0;
                tiling.set_core_rows(core_rows);
                tiling.set_core_rows_remain(core_rows_remain);
            }
        }
    }      
    tiling.set_aivNum(aivNum);
    
    std::cout << "B_length:" << tiling.get_B_length() << std::endl;
    std::cout << "T_length:" << tiling.get_T_length() << std::endl;
    std::cout << "C_length:" << tiling.get_C_length() << std::endl;
    std::cout << "tiling_size:" << tiling.get_tiling_size() << std::endl;
    std::cout << "core_rows:" << tiling.get_core_rows() << std::endl;
    std::cout << "core_rows_remain:" << tiling.get_core_rows_remain() << std::endl;
    std::cout << "aivNum:" << tiling.get_aivNum() << std::endl;
    std::cout << "X_length_align:" << tiling.get_X_length_align() << std::endl; 

    std::cout << "tilingkey:" << context->GetTilingKey() << std::endl;

    context->SetBlockDim(aivNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x_shape = context->GetInputShape(0);
    const gert::Shape* h0_shape = context->GetInputShape(2);
    gert::Shape* xr_shape = context->GetOutputShape(0);
    gert::Shape* xw_shape = context->GetOutputShape(1);
    gert::Shape* xk_shape = context->GetOutputShape(2);
    gert::Shape* xv_shape = context->GetOutputShape(3);
    gert::Shape* xa_shape = context->GetOutputShape(4);
    gert::Shape* xg_shape = context->GetOutputShape(5);
    gert::Shape* ht_shape = context->GetOutputShape(6);
    *xr_shape = *x_shape;
    *xw_shape = *x_shape;
    *xk_shape = *x_shape;
    *xv_shape = *x_shape;
    *xa_shape = *x_shape;
    *xg_shape = *x_shape;
    *ht_shape = *h0_shape;
    return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    auto data_type = context->GetInputDataType(0);
    context->SetOutputDataType(0, data_type);
    context->SetOutputDataType(1, data_type);
    context->SetOutputDataType(2, data_type);
    context->SetOutputDataType(3, data_type);
    context->SetOutputDataType(4, data_type);
    context->SetOutputDataType(5, data_type);
    context->SetOutputDataType(6, data_type);
    return GRAPH_SUCCESS;
}
}


namespace ops {
class TokenShift : public OpDef {
public:
    explicit TokenShift(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("rwkvag")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("h0")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("xr")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("xw")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("xk")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("xv")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("xa")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("xg")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("ht")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);
        this->SetInferDataType(ge::InferDataType);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b").AddConfig("ascend310p");

    }
};

OP_ADD(TokenShift);
}
